#!/usr/bin/env python

# Imports:
# -------
from config import *
import roslib; roslib.load_manifest(PKG)
import rosbag, log, wx, cv, cv_bridge
from StringIO import StringIO
from gui.main_frame import *

# ----------------
class BagLabeler():
# ----------------
	""" The guts of the bag labeling tool """

	# ----------------
	def __init__(self):
	# ----------------
		log.info('Initializing the bag labeler')

		# Initialize everything
		self.bagPath = None
		self.bag = None

		# Create our CV bridge
		self.bridge = cv_bridge.CvBridge()

	# -----------------
	def unloadBag(self):
	# -----------------
		if self.bag != None:
			log.info('Started unloading bag: %s' % self.bagPath)
			self.bag.close()
			log.info('Finished unloading bag: %s' % self.bagPath)

	# ---------------------------------------
	def loadBag(self, bagPath, callback=None):
	# ---------------------------------------
		# Unload the current bag (if we have one)
		self.unloadBag()

		# Store the bag path
		self.bagPath = bagPath
		log.info('Started loading ROS bag: %s' % bagPath)

		# Load the bag
		self.bag = rosbag.Bag(bagPath, 'r')

		# Pull out information about the topics in the bag
		topics = sorted(set([c.topic for c in self.bag._get_connections()]))
		self.topicMessageCounts = {}
		self.topicMessageTypes = {}
		for topic in topics:
			connections = list(self.bag._get_connections(topic))
			stamps = [entry.time.to_sec() for entry in self.bag._get_entries(connections)]
			
			self.topicMessageCounts[topic] = len(stamps)
			self.topicMessageTypes[topic] = connections[0].datatype

		# Filter the topic message for images & annotations
		self.imageTopics = []
		self.annotationTopics = []
		for topic, topicType in self.topicMessageTypes.items():
			if topicType == 'sensor_msgs/Image':
				# See if we can convert the image from the topic... 
				imageGenerator = self.bag.read_messages(topics=[topic])
				__, msg, __ = imageGenerator.next()
				if self.convertToBitmap(msg) != None:
					self.imageTopics.append(topic)
				else:
					log.warning('Unable to convert images from topic: %s, topic will be omitted' % topic)

			elif topicType == 'people_dataset_msgs/Annotation':
				self.annotationTopics.append(topic)

		# Set the image topic to be the first topic in the list
		self.setImageTopic(self.imageTopics[0])

		# Execute the callback
		if callback != None:
			callback(self.imageTopics, self.annotationTopics)
		log.info('Finished loading ROS bag: %s' % bagPath)

	# -------------------------------------------------
	def setImageTopic(self, imageTopic, startTime=None):
	# -------------------------------------------------
		self.imageTopic = imageTopic

		# Get all the image times
		self.msgTimes = [out[2] for out in self.bag.read_messages(topics = [imageTopic])]

		# Set the current time & index
		startIndex = 0
		if startTime != None:
			for i, time in enumerate(msgTimes):
				if startTime < time:
					startIndex = i
					break
		self.index = startIndex

	# -------------------------------
	def convertToBitmap(self, rosImg):
	# -------------------------------
		try:
			# Convert the image to something WX can display
			width =  rosImg.width
			height = rosImg.height
			cvImg = self.bridge.imgmsg_to_cv(rosImg, 'rgb8')
			wxImage = wx.EmptyImage(width, height)
			wxImage.SetData(cvImg.tostring())
			return wxImage.ConvertToBitmap()

		except AttributeError as e: 
			return None
			

	# -----------------------
	def getImage(self, index):
	# -----------------------
		# Get the image
		startTime = self.msgTimes[index]
		gen = self.bag.read_messages(topics=[self.imageTopic], start_time=startTime)
		topic, msg, t = gen.next()
		bmp = self.convertToBitmap(msg)		

		return topic, bmp, t

	# ----------------
	def getIndex(self):
	# ----------------
		return (self.index, len(self.msgTimes))

	# ------------
	def curr(self):
	# ------------
		return self.getImage(self.index)

	# ------------
	def next(self):
	# ------------		
		# Increment the index
		self.index += 1
		if self.index >= len(self.msgTimes):
			self.index = 0

		return self.getImage(self.index)

	# ------------
	def prev(self):
	# ------------
		# Decrement the index
		self.index -= 1
		if self.index <= 0:
			self.index = len(self.msgTimes) - 1

		return self.getImage(self.index)

	# ------------------
	def isBagLoaded(self):
	# ------------------
		return self.bag != None

	# ----------------
	def shutdown(self):
	# ----------------
		log.info('Shutting down the bag labeler')
		self.unloadBag()

# ---------------------
if __name__=='__main__':
# ---------------------
	# Create the labeler
	bagLabeler = BagLabeler()

	# Create the GUI
	app = wx.App()
	frame = MainFrame(bagLabeler, WINDOW_TITLE, WINDOW_SIZE)
	frame.Show()
	app.MainLoop()
